# -*- coding: utf-8 -*-
import geopandas as gpd
from shapely.geometry import box
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrow, Rectangle
import os

matplotlib.use("Agg")  # 避免 GUI 错误


def generate_a4_maps(in_fc, target_names, name_field="标准名称", output_dir=r"D:\Windows",
                     a4_width=8.27, a4_height=11.69, border_ratio=0.02, dpi=300,
                     max_batch_height_ratio=0.9):
    """
    输入Shapefile路径 + 道路列表，生成裁剪后的Shapefile + A4 PNG。
    如果道路太多，自动分批，每个文件只显示能够完整显示的道路。

    参数:
        in_fc: 原始Shapefile路径
        target_names: 需要显示的道路列表
        name_field: 道路名称字段
        output_dir: 输出目录
        a4_width, a4_height: A4画布尺寸(英寸)
        border_ratio: 黑色内边框比例
        dpi: 图片分辨率
        max_batch_height_ratio: A4纵向显示比例阈值（留空白边距）
    """
    os.makedirs(output_dir, exist_ok=True)
    gdf = gpd.read_file(in_fc)

    # 筛选目标道路
    highlight_all = gdf[gdf[name_field].isin(target_names)]

    # 按纵向顺序排序（ymin）
    highlight_all = highlight_all.assign(ymin=highlight_all.geometry.bounds['miny'])
    highlight_all = highlight_all.sort_values('ymin')

    # 分批存放道路
    batches = []
    current_batch = []
    current_ymin, current_ymax = None, None

    for idx, row in highlight_all.iterrows():
        road_bounds = row.geometry.bounds
        road_miny, road_maxy = road_bounds[1], road_bounds[3]

        if not current_batch:
            current_batch = [row]
            current_ymin, current_ymax = road_miny, road_maxy
        else:
            batch_height = road_maxy - current_ymin
            total_height = current_ymax - current_ymin
            if total_height <= max_batch_height_ratio * (total_height + (road_maxy - road_miny)):
                current_batch.append(row)
                current_ymin = min(current_ymin, road_miny)
                current_ymax = max(current_ymax, road_maxy)
            else:
                batches.append(current_batch)
                current_batch = [row]
                current_ymin, current_ymax = road_miny, road_maxy
    if current_batch:
        batches.append(current_batch)

    # 生成每个批次的文件
    for i, batch_rows in enumerate(batches):
        batch_gdf = gpd.GeoDataFrame(batch_rows, crs=gdf.crs)
        # 计算边界
        xmin, ymin, xmax, ymax = batch_gdf.total_bounds
        dx = (xmax - xmin) * border_ratio
        dy = (ymax - ymin) * border_ratio
        x0, y0 = xmin - dx, ymin - dy
        x1, y1 = xmax + dx, ymax + dy
        bbox = box(x0, y0, x1, y1)

        # 裁剪
        clipped = batch_gdf.clip(bbox)

        # 导出Shapefile
        safe_names = [r[name_field].replace("/", "_").replace("\\", "_") for r in batch_rows]
        shp_filename = f"batch{i + 1}_" + "_".join(safe_names) + ".shp"
        shp_path = os.path.join(output_dir, shp_filename)
        clipped.to_file(shp_path, encoding="utf-8")
        print(f"✅ 已生成 Shapefile: {shp_path}")

        # 绘制A4地图
        fig, ax = plt.subplots(figsize=(a4_width, a4_height))
        # 黑色内边框
        ax.add_patch(Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, edgecolor='black', linewidth=2))
        # 底图
        gdf.plot(ax=ax, color="lightgray", linewidth=0.5)
        # 高亮道路
        batch_gdf.plot(ax=ax, color="red", linewidth=2)
        # 道路标注
        for idx2, row2 in batch_gdf.iterrows():
            x, y = row2.geometry.representative_point().coords[0]
            ax.text(x, y, row2[name_field], fontsize=10, color="blue", ha="center", fontproperties="SimHei")
        ax.set_xlim(x0, x1)
        ax.set_ylim(y0, y1)
        ax.set_axis_off()
        # 指北针
        arrow_x = x1 - (x1 - x0) * 0.05
        arrow_y = y1 - (y1 - y0) * 0.05
        ax.add_patch(FancyArrow(arrow_x, arrow_y, 0, (y1 - y0) * 0.05,
                                width=(x1 - x0) * 0.002, color='black'))
        ax.text(arrow_x, arrow_y + (y1 - y0) * 0.055, "N",
                fontsize=12, fontweight='bold', ha='center')
        # 保存图像
        image_filename = f"batch{i + 1}_" + "_".join(safe_names) + "_A4_border.png"
        image_path = os.path.join(output_dir, image_filename)
        plt.savefig(image_path, dpi=dpi)
        plt.close()
        print(f"✅ 已生成 A4 地图图像（带内边框）: {image_path}")


# ================== 使用示例 ==================
if __name__ == "__main__":
    excelFileName = "南塘镇道路"
    in_fc = rf"D:\Data\南塘镇\{excelFileName}.shp"
    target_names = ["南胜大街", "南塘圩一街", "南大路"]  # 替换为你需要的道路
    output_dir = r"D:\Windows"

    generate_a4_maps(in_fc, target_names, name_field="标准名称", output_dir=output_dir)
